/**
 * @file ebex_bootstrap.c
 *
 * @date Oct 25, 2011
 * @author seth
 *
 * @brief This file is part of EBEX State Estimator, created for the EBEX project
 *
 * This software is copyright (C) 2011 Columbia University
 *
 * EBEX State Estimator is free software; you can redistribute it and/or modify
 * it under the terms of the GNU General Public License as published by
 * the Free Software Foundation; either version 2 of the License, or
 * (at your option) any later version.
 *
 * EBEX State Estimator is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License
 * along with EBEX State Estimator; if not, write to the Free Software Foundation, Inc.,
 * 59 Temple Place, Suite 330, Boston, MA  02111-1307  USA
 *
 */

#include <assert.h>
#include <memory_util.h>
#include <log_util.h>
#include <random_util.h>

#include <ebex_genetic.h>
#include <ebex_genetic_chromosome.h>
#include <ebex_de.h>
#include <gyros.h>
#include <ebex_quaternion.h>
#include <ebex_matrix.h>
#include <ebex_filter_cfg.h>
#include <ebex_model.h>
#include <ebex_bootstrap.h>
#include <ebex_genetic_stats.h>
#include <ebex_gyro_sc.h>
#include <dirfile_data.h>

static bool ebex_bootstrap_organism_scatter(population_t *m_pop, size_t m_chromosome)
{
	size_t i;
	ebex_model_t *model = (ebex_model_t*) m_pop->data;

	/// We record the scatter fitness at the end of an entire scatter cycle (all chromosomes)
	if (!m_chromosome)
	{
		/// If we are not improving, then end the cycle
		if ((m_pop->de_params->scatter_fittness > GENETIC_MIN_FITNESS) &&
				(m_pop->sorted_organisms[0]->fitness - m_pop->de_params->scatter_fittness) < model->parameters->bootstrap.de_scatter_min_improvement)
			return false;

		m_pop->de_params->scatter_fittness = m_pop->sorted_organisms[0]->fitness;
	}

	elog(log_level_verbose, "Scattering Chromosome %zu!", m_chromosome);

	/**
	 * Scattering leaves the best organism in place in the population but re-seeds all others.  This ensures
	 * that we never decrease our fitness but allows the solution to adapt to improved conditions in other
	 * chromosomes
	 */
#	pragma omp parallel for \
	   shared(m_pop) private(i) schedule(static)
	for (i = 0; i < m_pop->size; i++)
	{
		double allele_mean;
		double range;
		double min_range;
		double max_range;
		random_sample_t rndsamp;
		organism_t *entity = m_pop->sorted_organisms[i];
		if (entity->chromosome[m_chromosome].rnd.type == random_dist_none) continue;

		for (size_t allele = 0; allele < entity->chromosome[m_chromosome].length; allele++)
		{
			switch (entity->chromosome[m_chromosome].type)
			{
				case CHROMOSOME_TYPE_DOUBLE:

					allele_mean = (entity->chromosome[m_chromosome].allele_max_double - entity->chromosome[m_chromosome].allele_min_double) / 2.0;

					range = 1.5 * fabs(allele_mean - entity->chromosome[m_chromosome].dbl_allele[allele]);
					min_range = CLAMP(entity->chromosome[m_chromosome].dbl_allele[allele] - range, entity->chromosome[m_chromosome].allele_min_double, entity->chromosome[m_chromosome].allele_max_double);
					max_range = CLAMP(entity->chromosome[m_chromosome].dbl_allele[allele] + range, entity->chromosome[m_chromosome].allele_min_double, entity->chromosome[m_chromosome].allele_max_double);

					entity->chromosome[m_chromosome].dbl_allele[allele] = ran_double_range_uniform(min_range, max_range);

					break;

				case CHROMOSOME_TYPE_QUATERNION:
					rndsamp.type = random_dist_gaussian;
					rndsamp.scale = 1.0;
					rndsamp.offset = 0.0;
					rndsamp.gaussian.mean = 0.0;
					rndsamp.gaussian.stddev = sqrt(M_PI - entity->fitness);
					ebex_quat_randomize_population(&entity->chromosome[m_chromosome].q_allele[allele], &rndsamp, 1);
					break;
			}
		}

		m_pop->evaluate(m_pop, entity);
	}
	return true;
}

static bool ebex_bootstrap_organism_seed(population_t *m_pop, organism_t *m_entity)
{

	ebex_model_t *model = (ebex_model_t*)m_pop->data;

	m_entity->fitness = GENETIC_MIN_FITNESS;

	for (size_t chromosome = 0; chromosome < num_gyro_params; chromosome++)
	{
		for (size_t j = 0; j < model->num_gyroboxes; j++)
		{
			for (size_t k = 0; k < 3; k++)
			{
				size_t offset = j * 3 + k;
				switch(m_pop->seed_strategy)
				{
					case GENETIC_SEED_GAUSSIAN:
						ebex_genetic_seed_dbl_rnd_unit_gaussian(m_pop, m_entity, chromosome, offset);
						break;
					case GENETIC_SEED_NONE:
						break;
					case GENETIC_SEED_SAMPLE:
						ebex_genetic_seed_dbl_rnd_sample(m_pop, m_entity, chromosome, offset);
						break;
					case GENETIC_SEED_UNIFORM:
						ebex_genetic_seed_dbl_rnd_uniform(m_pop, m_entity, chromosome, offset);
						break;
					case GENETIC_SEED_ZERO:
						ebex_genetic_seed_dbl_zero(m_pop, m_entity, chromosome, offset);
						break;
					default:
						die("Unknown seeding strategy");
						break;
				}
			}
		}
	}

	for (size_t i = 0; i < model->num_gyroboxes; i++)
	{
		size_t gyrobox_chromosome = i + num_gyro_params;
		switch(m_pop->seed_strategy)
		{
			case GENETIC_SEED_SAMPLE:
			case GENETIC_SEED_GAUSSIAN:
				ebex_genetic_seed_quat_sample(m_pop, m_entity, gyrobox_chromosome, 0);
				break;
			case GENETIC_SEED_NONE:
				break;
			case GENETIC_SEED_UNIFORM:
				ebex_genetic_seed_quat_random(m_pop, m_entity, gyrobox_chromosome, 0);
				break;
			case GENETIC_SEED_ZERO:
				ebex_genetic_seed_quat_zero(m_pop, m_entity, gyrobox_chromosome, 0);
				break;
			default:
				die("Unknown seeding strategy");
				break;
		}

		/// The reference Gyrobox is always at the origin
		if (i == model->parameters->reference_gyrobox) ebex_genetic_seed_quat_zero(m_pop, m_entity, gyrobox_chromosome, 0);
	}

	for (size_t i = 0; i < model->num_sc; i++)
	{
		size_t sc_chromosome = i + num_gyro_params + model->num_gyroboxes;
		switch(m_pop->seed_strategy)
		{
			case GENETIC_SEED_SAMPLE:
			case GENETIC_SEED_GAUSSIAN:
				ebex_genetic_seed_quat_sample(m_pop, m_entity, sc_chromosome, 0);
				break;
			case GENETIC_SEED_NONE:
				break;
			case GENETIC_SEED_UNIFORM:
				ebex_genetic_seed_quat_random(m_pop, m_entity, sc_chromosome, 0);
				break;
			case GENETIC_SEED_ZERO:
				ebex_genetic_seed_quat_zero(m_pop, m_entity, sc_chromosome, 0);
				break;
			default:
				die("Unknown seeding strategy");
				break;
		}
	}

	return true;
}

static void ebex_bootstrap_scale_rotate_omega(ebex_3vec_t *m_output_scale, population_t *m_pop, gyrobox_t *m_gyrobox, ebex_gyro_reading_t *m_gyrodata, size_t m_segment)
{
	ebex_model_t *model = (ebex_model_t*)m_pop->data;
	organism_t *best_entity = m_pop->sorted_organisms[0];

	ebex_quat_t attitude;
	ebex_quat_t end_attitude;
	ebex_quat_t err_quat;

	/// Rotate the beginning star camera reading into the gyro frame
	ebex_quat_mul(&attitude, &model->data->sc_segment[m_segment].attitude, best_entity->chromosome[num_gyro_params + model->num_gyroboxes + model->data->sc_segment[m_segment].sc_index].q_allele);
	/// Rotate the end star camera reading into the gyro frame
	ebex_quat_mul(&end_attitude, &model->data->sc_segment[m_segment+1].attitude, best_entity->chromosome[num_gyro_params + model->num_gyroboxes + model->data->sc_segment[m_segment+1].sc_index].q_allele);

	ebex_gyro_sc_quat_propagate_forward(model, m_gyrodata, &attitude, m_gyrobox, NULL, 0, NULL, NULL, m_segment);

	ebex_quat_error(&err_quat, &attitude, &end_attitude);
	ebex_3vec_euler_from_quat(m_output_scale, rotation_order_XYZ, &err_quat);
	ebex_3vec_scalar_div(m_output_scale, m_output_scale, -2.0 * (model->data->sc_segment[m_segment + 1].time - model->data->sc_segment[m_segment].time));
}

/**
 * Takes the best individual in the bootstrap population and outputs calculated state
 * information (yaw, pitch, roll, q1-4, etc)
 * @param m_pop Population from which to output state
 * @return True on success, false on failure
 */
static bool ebex_bootstrap_output(population_t *m_pop)
{
	ebex_model_t *model = (ebex_model_t*)m_pop->data;

	ebex_quat_t start_attitude;
	ebex_quat_t attitude;
	ebex_quat_t end_attitude;

	size_t quat_vec_size = 0;
	ebex_quat_t *output_quat_vec[2] = {NULL};
	double		*output_time_vec = NULL;

	size_t box;
	size_t i;
	gyrobox_t gyrobox = {0};

	organism_t *best_entity = m_pop->sorted_organisms[0];

	if (!model->parameters->bootstrap.output) return false;

	for (box = 0; box < model->num_gyroboxes; box++)
	{
		/// Copy the gyro parameters from the entity chromosome into a useable, local format
		for (size_t j = 0; j < 3; j++)
		{
			memcpy(&gyrobox.gyro[j].param, &model->gyrobox[box].gyro[j].param, sizeof(gyro_param_t));
			for (i = 0; i < num_gyro_params; i++)
			{
				size_t offset = box * 3 + j;
				gyrobox.gyro[j].param.value.array[i] = best_entity->chromosome[i].dbl_allele[offset];
			}

			gyrobox.gyro[j].param.sample.rate_noise.type = random_dist_none;

		}
		ebex_quat_copy(&gyrobox.alignment, best_entity->chromosome[box + num_gyro_params].q_allele);


		for (i = 0; i < model->data->num_segments; i++)
		{
			ebex_gyro_reading_t *gyrodata;
			size_t steps;
			ebex_3vec_t offset = EBEX_3VEC_ZERO;
			ebex_quat_t *sc_rot = best_entity->chromosome[num_gyro_params + model->num_gyroboxes + model->data->sc_segment[i].sc_index].q_allele;
			ebex_quat_t sc_inv;


			gyrodata = &model->data->gyro_segment[box][i];
			if (model->parameters->bootstrap.offset_solution)
				ebex_bootstrap_scale_rotate_omega(&offset, m_pop, &gyrobox, gyrodata, i);

			/// Rotate the beginning star camera reading into the gyro frame
			ebex_quat_mul(&start_attitude, &model->data->sc_segment[i].attitude, sc_rot);
			ebex_quat_mul(&end_attitude, &model->data->sc_segment[i + 1].attitude, sc_rot);
			ebex_quat_inv(&sc_inv, sc_rot);

			ebex_quat_copy(&attitude, &start_attitude);

			if (quat_vec_size < gyrodata->num_obs)
			{
				quat_vec_size = 2 * gyrodata->num_obs;
				output_quat_vec[0] = e_realloc(output_quat_vec[0], 2 * quat_vec_size * sizeof(ebex_quat_t));
				output_quat_vec[1] = output_quat_vec[0] + quat_vec_size;
				output_time_vec = e_realloc(output_time_vec, quat_vec_size * sizeof(double));
			}
			e_memset(output_quat_vec[0], 0, 2 * quat_vec_size * sizeof(ebex_quat_t));
			steps = ebex_gyro_sc_quat_propagate_forward(model, gyrodata, &attitude, &gyrobox, output_quat_vec[0], 1, &offset, NULL, i);
			if (model->parameters->bootstrap.backward_smoothing)
			{
				ebex_quat_copy(&attitude, &end_attitude);
				ebex_gyro_sc_quat_propagate_backward(model, gyrodata, &attitude, &gyrobox, output_quat_vec[1], 1, NULL, NULL, i, steps);
				for (size_t j = 0; j < steps; j++)
				{
					ebex_quat_slerp(&output_quat_vec[0][j], (double)(j)/ (double)steps, &output_quat_vec[0][j], &output_quat_vec[1][j]);
				}
			}
			ebex_gyro_sc_output_time(model->parameters->bootstrap.output, gyrodata->time, steps);
			ebex_gyro_sc_output_attitude(model->parameters->bootstrap.output, output_quat_vec[0], NULL, &sc_inv, steps);
			model->parameters->bootstrap.output->buffer_produced += steps;

			ebex_data_write_buffer(model->parameters->bootstrap.output);

		}
	}

	e_free_quiet(output_quat_vec[0]);
	e_free_quiet(output_time_vec);
	return true;
}

static bool ebex_bootstrap_organism_evaluate(population_t *m_pop, organism_t *m_entity)
{
	ebex_model_t *model = (ebex_model_t*)m_pop->data;

	double mse = 0.0;
	double penalty[3] = {0.0};
	double total_penalty = 0.0;
	double mean;

	ebex_quat_t start_attitude;
	ebex_quat_t attitude;
	ebex_quat_t end_attitude;
	ebex_quat_t err_quat;

	size_t box;
	size_t i;
	double total_time = 0.0;
	gyrobox_t gyrobox = {0};


	for (box = 0; box < model->num_gyroboxes; box++)
	{
		for (size_t j = 0; j < 3; j++)
		{
			memcpy(&gyrobox.gyro[j].param, &model->gyrobox[box].gyro[j].param, sizeof(gyro_param_t));
			for (i = 0; i < num_gyro_params; i++)
			{
				size_t offset = box * 3 + j;
				gyrobox.gyro[j].param.value.array[i] = m_entity->chromosome[i].dbl_allele[offset];
			}
			/// Prevent the propagation routine from adding artificial noise
			gyrobox.gyro[j].param.sample.rate_noise.type = random_dist_none;


			mean = (gyrobox.gyro[j].param.max.misalignment - gyrobox.gyro[j].param.min.misalignment) / 2.0;
			penalty[j] = (mean - gyrobox.gyro[j].param.value.misalignment) / mean;
		}
		total_penalty += sqrt((pow(penalty[0], 2.0) + pow(penalty[1], 2.0) + pow(penalty[2], 2.0)) / 3.0);
		ebex_quat_copy(&gyrobox.alignment, m_entity->chromosome[box + num_gyro_params].q_allele);

#		pragma omp parallel for \
			shared(m_pop, model, mse, total_time) private(attitude, err_quat, start_attitude, end_attitude, i) schedule(dynamic)
		for (i = 0; i < model->data->used_segments; i++)
		{
			double radians_err;

			/// Rotate the beginning star camera reading into the gyro frame
			ebex_quat_mul(&start_attitude, &model->data->sc_segment[i].attitude, m_entity->chromosome[num_gyro_params + model->num_gyroboxes + model->data->sc_segment[i].sc_index].q_allele);
			/// Rotate the end star camera reading into the gyro frame
			ebex_quat_mul(&end_attitude, &model->data->sc_segment[i+1].attitude, m_entity->chromosome[num_gyro_params + model->num_gyroboxes + model->data->sc_segment[i+1].sc_index].q_allele);

			ebex_quat_copy(&attitude, &start_attitude);

			ebex_gyro_sc_quat_propagate_forward(model, &model->data->gyro_segment[box][i], &attitude, &gyrobox, NULL, 0, NULL, NULL, i);


			ebex_quat_error(&err_quat, &attitude, &end_attitude);
			if (isnan(err_quat.w))
			{
				elog(log_level_debug, "got NaN!");
			}
			radians_err = 2.0 * acos(err_quat.w);
//			if (M_PI - m_pop->sorted_organisms[0]->fitness < 49.0 * M_PI / 3600.0 / 180.0 )
//				elog(log_level_verbose, "Segment %zu, error %g", i, (3600.0 * 180.0 / M_PI) * radians_err);
			mse += (radians_err * radians_err * (model->data->sc_segment[i + 1].time - model->data->sc_segment[i].time));
			total_time += (model->data->sc_segment[i + 1].time - model->data->sc_segment[i].time);
		}
	}

	m_entity->fitness = (M_PI -  sqrt(mse / total_time)) ;

	if (m_entity->fitness < GENETIC_MIN_FITNESS) m_entity->fitness = GENETIC_MIN_FITNESS;

	return true;
}

static bool ebex_bootstrap_post_gen_hook(const size_t m_count, population_t *m_pop, organism_t *m_entity __attribute__((unused)))
{
	ebex_model_t *model = (ebex_model_t *)m_pop->data;

	double min_error = (M_PI - m_pop->sorted_organisms[0]->fitness) * 180.0 / M_PI;
	double max_error = (M_PI - m_pop->sorted_organisms[m_pop->size-1]->fitness) * 180.0 / M_PI;

	if (min_error < 1.0 / 60.0 && max_error < 1.0 / 6.0)
	{
		elog(log_level_verbose,
				"After generation %zu, population has errors between %0.2f\" and %0.2f\"",
				m_count, min_error * 3600.0 , max_error * 3600.0 );
	}
	else if (min_error < 1.0 && max_error < 1.0)
	{
		elog(log_level_verbose,
				"After generation %zu, population has errors between %0.2f\' and %0.2f\'",
				m_count, min_error * 60.0 , max_error * 60.0 );
	}
	else
	{
		elog(log_level_verbose,
				"After generation %zu, population has errors between %0.2f degrees and %0.2f degrees",
				m_count, min_error , max_error );
	}

	if (min_error < model->parameters->bootstrap.err_tol)
	{
		elog(log_level_normal, "Reached desired error of %g degrees", model->parameters->bootstrap.err_tol);
		return false;
	}

	if (max_error - min_error < model->parameters->bootstrap.var_tol)
	{
		elog(log_level_normal, "Minimum variability of %g degrees reached in population", model->parameters->bootstrap.var_tol);
		if (m_pop->de_params->scatter)
		{
			if (m_pop->de_params->scatter_chromosome > m_pop->num_chromosomes) m_pop->de_params->scatter_chromosome = 0;

			return ebex_bootstrap_organism_scatter(m_pop, m_pop->de_params->scatter_chromosome++);
		}

		return false;
	}

	return true;
}

static bool ebex_bootstrap_pre_gen_hook(const size_t m_count, population_t *m_pop, organism_t *m_entity __attribute__((unused)))
{

	if (shutdown_now)
	{
		elog(log_level_normal, "Received shutdown signal.  Exiting bootstrap filter.");
		return false;
	}

	elog(log_level_verbose, "Population size is %zu at start of generation %zu", m_pop->orig_size, m_count);

	return true;
}

static void ebex_bootstrap_copy_best_entity(population_t *m_src, population_t *m_dest)
{
	organism_t *best_entity = m_src->sorted_organisms[0];
	organism_t *storage;

	storage = genetic_get_new_organism(m_dest);

	genetic_organism_copy(m_dest, storage, best_entity);
}

static void ebex_bootstrap_store_parameters(population_t *m_pop)
{
	ebex_model_t *model = m_pop->data;
	organism_t *best_entity = m_pop->sorted_organisms[0];

	for (size_t chromosome = 0; chromosome < num_gyro_params; chromosome++)
	{
		for (size_t j = 0; j < model->num_gyroboxes; j++)
		{
			for (size_t k = 0; k < 3; k++)
			{
				size_t offset = j * 3 + k;
				model->gyrobox[j].gyro[k].param.value.array[chromosome] = best_entity->chromosome[chromosome].dbl_allele[offset];
			}
		}
	}

	for (size_t i = 0; i < model->num_gyroboxes; i++)
	{
		size_t gyrobox_chromosome = i + num_gyro_params;
		ebex_quat_copy(&model->gyrobox[i].alignment, &best_entity->chromosome[gyrobox_chromosome].q_allele[i]);
	}

	for (size_t i = 0; i < model->num_sc; i++)
	{
		size_t sc_chromosome = i + num_gyro_params + model->num_gyroboxes;
		ebex_quat_copy(&model->sc[i].alignment, &best_entity->chromosome[sc_chromosome].q_allele[i]);
	}
}

void ebex_state_bootstrap(ebex_model_t *m_model, population_t *m_storage)
{
	population_t *bootstrap_pop;

	bootstrap_pop = genetic_pop_new(m_model->parameters->bootstrap.population_size, m_model->parameters->num_gyros + m_model->parameters->num_sc + num_gyro_params );
	ebex_diffevolv_setparameters(bootstrap_pop,
			m_model->parameters->bootstrap.de_strategy,
			m_model->parameters->bootstrap.de_crossover,
			m_model->parameters->bootstrap.de_numvec,
			m_model->parameters->bootstrap.de_minweight,
			m_model->parameters->bootstrap.de_maxweight,
			m_model->parameters->bootstrap.de_crossover_ratio,
			m_model->parameters->bootstrap.de_opposition_rate,
			m_model->parameters->bootstrap.de_reference,
			m_model->parameters->bootstrap.de_scatter);

	bootstrap_pop->seed_strategy = m_model->parameters->bootstrap.seed_strategy;
	bootstrap_pop->seed = ebex_bootstrap_organism_seed;
	bootstrap_pop->chromosome_constructor = ebex_gyro_sc_generate_chromosome;
	bootstrap_pop->chromosome_destructor = ebex_genetic_chromosome_free;
	bootstrap_pop->verify = ebex_gyro_sc_organism_verify;
	bootstrap_pop->evaluate = ebex_bootstrap_organism_evaluate;

	bootstrap_pop->post_generation_hook = ebex_bootstrap_post_gen_hook;
	bootstrap_pop->pre_generation_hook = ebex_bootstrap_pre_gen_hook;


	m_model->private_data = bootstrap_pop;
	bootstrap_pop->data = m_model;

	while (ebex_differential_evolution(bootstrap_pop, m_model->parameters->bootstrap.max_iter) >= m_model->parameters->bootstrap.max_iter
			&& !m_model->parameters->bootstrap.force_exit)
	{
		genetic_genocide(bootstrap_pop, 0);
	}

	ebex_state_print_stats(bootstrap_pop);

	ebex_gyro_sc_register_attitude_output(m_model, &m_model->parameters->bootstrap.output, m_model->parameters->bootstrap.output_filename, m_model->parameters->bootstrap.output_type);
	ebex_bootstrap_output(bootstrap_pop);
	ebex_data_close_source(m_model->parameters->bootstrap.output);

	if (m_storage) ebex_bootstrap_copy_best_entity(bootstrap_pop, m_storage);
	else ebex_bootstrap_store_parameters(bootstrap_pop);

	bootstrap_pop->data = NULL;
	genetic_close(bootstrap_pop);
}
